<?php

/**
 * This file is part of Webman AI.
 *
 * @author    walkor<walkor@workerman.net>
 * @copyright walkor<walkor@workerman.net>
 * @link      http://www.workerman.net/
 */

namespace plugin\ai\app\controller;

use Exception;
use plugin\ai\app\event\data\ModelRequestData;
use plugin\ai\app\event\data\ModelResponseData;
use plugin\ai\app\model\AiDataset;
use plugin\ai\app\model\AiEmbedding;
use plugin\ai\app\modelHandler\Embedding;
use support\Redis;
use support\Request;
use support\Response;
use Throwable;
use Webman\Event\Event;
use Workerman\Protocols\Http\Chunk;

/**
 * Embedding消息
 */
class EmbeddingController extends Base
{

    /**
     * 不需要登录的方法
     *
     * @var string[]
     */
    protected $noNeedLogin = ['search', 'list'];

    /**
     * 发消息
     * @param Request $request
     * @return Response
     * @throws Exception|Throwable
     */
    public function search(Request $request): ?Response
    {
        if (!class_exists(\Redis::class)) {
            return json(['error' => ['message' => '请安装redis扩展']]);
        }
        // chatgpt 配置
        $content = $request->post('content');
        $connection = $request->connection;
        $dataset = $request->post('dataset');
        $datasetInfo = AiDataset::find($dataset);
        if (!$datasetInfo) {
            return json(['error' => ['message' => '数据集不存在']]);
        }
        $model = $datasetInfo->model;

        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'input' => $content,
            'encodding_format' => 'float'
        ];
        $modelRequestData->options = [
            'complete' => function ($data) use ($connection, $model, $dataset, $modelRequestData) {
                $responseData = new ModelResponseData();
                $responseData->data = $data;
                $responseData->modelRequestData = $modelRequestData;
                Event::dispatch('ai.embedding.create.response', $responseData);
                $data = $responseData->data;
                $buffer = json_encode($data);
                if (isset($data['error'])) {
                    $connection->send(new Chunk($buffer));
                    $connection->send(new Chunk(''));
                    return;
                }
                $embedding = $data['data'][0]['embedding'];
                $blob = '';
                foreach ($embedding as $value) {
                    $blob .= pack('f', $value);
                }
                $index = "ai-embedding-$dataset";
                $count = 4;

                try {
                    try {
                        $indexExist = Redis::connection('plugin.ai.default')->rawCommand('FT.INFO', $index);
                    } catch (Throwable $e) {
                        $indexExist = false;
                    }
                    $res = [];
                    if ($indexExist) {
                        $res = Redis::connection('plugin.ai.default')->rawCommand('FT.SEARCH', $index, '*=>[KNN ' . $count . ' @text_embedding $blob]', 'PARAMS', '2', 'blob', $blob, 'SORTBY', '__text_embedding_score', 'DIALECT', '2');
                    }
                } catch (Exception $e) {
                    $connection->send(new Chunk(json_encode(['error' => ['message' => $e->getMessage()]])));
                    $connection->send(new Chunk(''));
                    return;
                }

                if (!isset($res[2][3])) {
                    $connection->send(new Chunk(json_encode(['code' => 0, 'msg' => 'ok', 'data' => []], JSON_PRETTY_PRINT | JSON_UNESCAPED_UNICODE)));
                    $connection->send(new Chunk(''));
                    return;
                }
                $embeddingIds = [];
                for ($i=0; $i < $count; $i++) {
                    if (empty($res[$i + 2][3])) {
                        continue;
                    }
                    $item = json_decode($res[$i + 2][3] , true);
                    if (!isset($item['id'])) {
                        continue;
                    }
                    $embeddingIds[] = $item['id'];
                }
                $searchText = [];
                $embeddings = AiEmbedding::select(['text', 'id'])->whereIn('id', $embeddingIds)->get()->keyBy('id');
                $embeddings = array_replace(array_flip($embeddingIds), $embeddings->toArray());
                foreach ($embeddings as $embedding) {
                    if (!isset($embedding['text'])) {
                        continue;
                    }
                    $searchText[] = $embedding['text'];
                }
                $connection->send(new Chunk(json_encode(['code' => 0, 'msg' => 'ok', 'data' =>$searchText], JSON_PRETTY_PRINT | JSON_UNESCAPED_UNICODE)));
                $connection->send(new Chunk(''));
            }
        ];
        Event::dispatch('ai.embedding.create.request', $modelRequestData);
        \plugin\ai\api\Embedding::create($modelRequestData->data, $modelRequestData->options);

        return response("\n")->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

}